package fem2.element;

import fem2.EASMode;
import fem2.EASModeManager;
import fem2.MaterialModel;
import fem2.MeshPart;
import fem2.enu.EASModeType;
import fem2.enu.MeshPartType;
import fem2.enu.StressType;
import inf.jlinalg.Array1DMatrix;
import inf.jlinalg.Array2DMatrix;
import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IMatrixReadOnly;
import inf.jlinalg.IVector;
import inf.jlinalg.IVectorReadOnly;
import inf.math.LinAlg;
import math2.LinAlgUtilities;

/**
 * enhanced assumed strain element based on paper of SImo & Rifai
 * 
 * TODO: this element give bad result for recover stress for the strip footing
 * example. Need check.
 * 
 * @author hbui
 * 
 */
public class EASStructuralElement extends StructuralElement {

	private double[] uHatLocal; // buffer uHatLocal
	private IVector alpha; // enhanced strain modes
	private EASMode mode;

	public EASStructuralElement(MeshPart mp, MaterialModel mm, EASModeType modeType) {
		super(mp, mm);

		/*
		 * check the compatibility of incompatible mode
		 */
		if (modeType == EASModeType.Q1E4 || modeType == EASModeType.Q1E5) {
			if ((mp.getType() != MeshPartType.QUAD4) && (mp.getType() != MeshPartType.QUAD8)) {
				throw new Error("the provided incompatible mode and the element is not compatible");
			}
		}

		this.mode = EASModeManager.getInstance().getMode(modeType);
	}

	@Override
	public IMatrixReadOnly computeKte(IMatrix ke) {
		MeshPart mp = getMeshPart();
		int n = mp.getBasis().getDimension();
		int dim = mp.getDimension();
		int nd = mode.countEnhancedStrainModes();
		int ntrs = dim * (dim + 1) / 2;
		/*
		 * initialize B
		 */
		IMatrix B = new Array2DMatrix(ntrs, dim * n);
		/*
		 * initialize G
		 */
		IMatrix G = new Array2DMatrix(ntrs, nd);
		/*
		 * initalize H
		 */
		Array1DMatrix H = Array1DMatrix.createRowMajor(nd, nd);
		/*
		 * initalize Gamma
		 */
		IMatrix Gamma = new Array2DMatrix(nd, dim * n);

		int ng = getNumIntegrationPoints();
		for (int k = 0; k < ng; k++) {
			double w = getIntegrationWeight(k);
			double[] xi = getIntegrationPoint(k);
			double[][] jac = mp.jacXAt(xi);
			double detJ = LinAlgUtilities.determinantNxN(jac);
			double coeff = w * detJ * getMaterialModel().getSectionProperty();
			IMatrix C = getMaterialModel().getConstitutiveMatrix(this, k);

			computeB(mp, B, xi);

			LinAlgUtilities.btCB(B, C, coeff, ke);

			mode.computeEnhancedStrainOperator(mp, G, xi);
			LinAlgUtilities.btCB(G, C, coeff, H);

			LinAlgUtilities.atCB(G, B, C, coeff, Gamma);
		}

		LinAlgUtilities.gtHiG(Gamma, H, -1.0, ke);
		return ke;
	}

	@Override
	public IVectorReadOnly computeRie(IVector rie) {
		MeshPart mp = getMeshPart();
		int n = mp.getBasis().getDimension();
		int dim = mp.getDimension();
		int nd = mode.countEnhancedStrainModes();
		int ntrs = dim * (dim + 1) / 2;
		/*
		 * initialize B
		 */
		IMatrix B = new Array2DMatrix(ntrs, dim * n);
		/*
		 * initialize G
		 */
		IMatrix G = new Array2DMatrix(ntrs, nd);
		/*
		 * initalize H
		 */
		Array1DMatrix H = Array1DMatrix.createRowMajor(nd, nd);
		/*
		 * initalize Gamma
		 */
		IMatrix Gamma = new Array2DMatrix(nd, dim * n);
		/*
		 * initialize h
		 */
		ArrayVector h = new ArrayVector(nd);

		int ng = getNumIntegrationPoints();
		for (int k = 0; k < ng; k++) {
			double w = getIntegrationWeight(k);
			double[] xi = getIntegrationPoint(k);
			double[][] jac = mp.jacXAt(xi);
			double detJ = LinAlgUtilities.determinantNxN(jac);
			double coeff = w * detJ * getMaterialModel().getSectionProperty();
			IVector s = getStress(k);
			IMatrix C = getMaterialModel().getConstitutiveMatrix(this, k);

			computeB(mp, B, xi);

			BLAM.multiply(coeff, BLAM.TRANSPOSE, B, s, 1.0, rie);

			mode.computeEnhancedStrainOperator(mp, G, xi);

			BLAM.multiply(coeff, BLAM.TRANSPOSE, G, s, 1.0, h);

			LinAlgUtilities.btCB(G, C, coeff, H);

			LinAlgUtilities.atCB(G, B, C, coeff, Gamma);
		}

		LinAlgUtilities.backSolve(H, h.getValues());
		BLAM.multiply(-1.0, BLAM.TRANSPOSE, Gamma, h, 1.0, rie);

		return rie;
	}

	@Override
	public void Initialize() {
		super.Initialize();
		/*
		 * initialize uHatLocal
		 */
		uHatLocal = new double[getMeshPart().countDofs()];
		/*
		 * initialize alpha
		 */
		alpha = new ArrayVector(mode.countEnhancedStrainModes());
	}

	@Override
	public void FinalizeNonlinearIteration() {
		getMaterialModel().FinalizeNonlinearIteration(this);
		/*
		 * 
		 */
		MeshPart mp = getMeshPart();
		int n = mp.getBasis().getDimension();
		int dim = mp.getDimension();
		int nd = mode.countEnhancedStrainModes();
		int ntrs = dim * (dim + 1) / 2;
		/*
		 * initialize B
		 */
		IMatrix B = new Array2DMatrix(ntrs, dim * n);
		/*
		 * initialize G
		 */
		IMatrix G = new Array2DMatrix(ntrs, nd);
		/*
		 * initalize H
		 */
		Array1DMatrix H = Array1DMatrix.createRowMajor(nd, nd);
		/*
		 * initalize Gamma
		 */
		IMatrix Gamma = new Array2DMatrix(nd, dim * n);
		/*
		 * initialize h
		 */
		ArrayVector h = new ArrayVector(nd);

		int ng = getNumIntegrationPoints();
		for (int k = 0; k < ng; k++) {
			double w = getIntegrationWeight(k);
			double[] xi = getIntegrationPoint(k);
			double[][] jac = mp.jacXAt(xi);
			double detJ = LinAlgUtilities.determinantNxN(jac);
			double coeff = w * detJ * getMaterialModel().getSectionProperty();
			IVector s = getStress(k);
			IMatrix C = getMaterialModel().getConstitutiveMatrix(this, k);

			computeB(mp, B, xi);

			mode.computeEnhancedStrainOperator(mp, G, xi);

			BLAM.multiply(coeff, BLAM.TRANSPOSE, G, s, 1.0, h);

			LinAlgUtilities.btCB(G, C, coeff, H);

			LinAlgUtilities.atCB(G, B, C, coeff, Gamma);
		}
		/*
		 * update alpha
		 */
		ArrayVector du = new ArrayVector(mp.getUHatLocal().clone());
		LinAlg.add(uHatLocal.length, -1.0, uHatLocal, 1.0, du.getValues());
		BLAM.multiply(1.0, BLAM.NO_TRANSPOSE, Gamma, du, -1.0, h);
		LinAlgUtilities.backSolve(H, h.getValues());
		BLAM.add(-1.0, h, alpha);
		/*
		 * update uHatLocal buffer
		 */
		LinAlg.copy(n, mp.getUHatLocal(), uHatLocal);
	}

	@Override
	public IVector computeStrain(double... xi) {
		MeshPart mp = getMeshPart();
		int n = mp.getBasis().getDimension();
		int dim = mp.getDimension();
		int nd = mode.countEnhancedStrainModes();
		int ntrs = dim * (dim + 1) / 2;
		IMatrix B = new Array2DMatrix(ntrs, dim * n);
		IMatrix G = new Array2DMatrix(ntrs, nd);
		IVector e = new ArrayVector(ntrs);
		/*
		 * compute standard strain
		 */
		computeB(mp, B, xi);
		IVector u = new ArrayVector(mp.getUHatLocal());
		BLAM.multiply(1.0, BLAM.NO_TRANSPOSE, B, u, 0.0, e);
		/*
		 * add the enhanced strain
		 */
		mode.computeEnhancedStrainOperator(mp, G, xi);
		BLAM.multiply(1.0, BLAM.NO_TRANSPOSE, G, alpha, 1.0, e);
		return e;
	}

	public IVector getRecoverStress(double... xi) {
		MeshPart mp = getMeshPart();
		int n = mp.getBasis().getDimension();
		int dim = mp.getDimension();
		int nd = mode.countRecoverStressModes();
		int ntrs = dim * (dim + 1) / 2;
		/*
		 * initialize B
		 */
		IMatrix B = new Array2DMatrix(ntrs, dim * n);
		/*
		 * initialize G
		 */
		IMatrix G = new Array2DMatrix(ntrs, nd);
		/*
		 * initalize H
		 */
		Array1DMatrix H = Array1DMatrix.createRowMajor(nd, nd);
		/*
		 * initalize Gamma
		 */
		IMatrix Gamma = new Array2DMatrix(nd, dim * n);
		/*
		 * initialize sigma
		 */
		ArrayVector sigma = new ArrayVector(ntrs);

		int ng = getNumIntegrationPoints();
		for (int k = 0; k < ng; k++) {
			double w = getIntegrationWeight(k);
			double[] xii = getIntegrationPoint(k);
			double[][] jac = mp.jacXAt(xi);
			double detJ = LinAlgUtilities.determinantNxN(jac);
			double t = getMaterialModel().getSectionProperty();
			double coeff = w * detJ * t;
			IMatrix C = getMaterialModel().getConstitutiveMatrix(this, k);

			computeB(mp, B, xii);

			mode.computeEnhancedStressOperator(mp, G, xii);

			LinAlgUtilities.gtHiG(G, (Array1DMatrix) C, coeff, H);

			BLAM.multiply(coeff, BLAM.TRANSPOSE, G, BLAM.NO_TRANSPOSE, B, 1.0, Gamma);
		}
		ArrayVector ue = new ArrayVector(mp.getUHatLocal());
		ArrayVector tmp = new ArrayVector(mode.countRecoverStressModes());
		BLAM.multiply(1.0, BLAM.NO_TRANSPOSE, Gamma, ue, 0.0, tmp);
		LinAlgUtilities.backSolve(H, tmp.getValues());
		mode.computeEnhancedStressOperator(mp, G, xi);
		BLAM.multiply(1.0, BLAM.NO_TRANSPOSE, G, tmp, 0.0, sigma);
		return sigma;
	}

	@Override
	public IVector getStress(int k, StressType type) {
		switch (type) {
		case RECOVER_STRESS_EAS:
			double[] xi = getIntegrationPoint(k);
			return getRecoverStress(xi);
		default:
			return super.getStress(k, type);
		}
	}
}
